﻿using System;
using System.Collections.Generic;
using System.Reflection;
using System.Runtime.CompilerServices;
using System.Reflection.Emit;

namespace ES.Interceptor
{
    internal sealed class MethodHook32 : IMethodHook
    {
        private MethodInfo _dynamicMethod;
        private IntPtr _originalMethodPointer;
        private IntPtr _hookMethodPointer;
        private UInt32 _savedOriginalMethodAddress;
        private Object _instance;      

        public void Initialize( MethodInfo origMethod, MethodInfo hookMethod, Object instance = null )
        {
            if ( origMethod == null ) throw new ArgumentNullException( "origMethod" );
            if ( hookMethod == null ) throw new ArgumentNullException( "hookMethod" );

            if ( !AreMethodsSignaturesEqual( origMethod, hookMethod ) )
            {
                throw new Exception( "Method doesn't have the same signature" );
            }

            // TODO: verify if it is really necessary to precompile methods
            RuntimeHelpers.PrepareMethod( origMethod.MethodHandle );
            RuntimeHelpers.PrepareMethod( hookMethod.MethodHandle );

            _instance = null;
            if ( !origMethod.IsStatic )
            {
                if ( instance == null )
                {
                    throw new ArgumentNullException( "You must specify the instance on which invoke the instance method" );
                }

                _instance = instance;
            }

            _originalMethodPointer = GetMethodAddress( origMethod );
            _hookMethodPointer = GetMethodAddress( hookMethod );

            unsafe
            {
                UInt32* ptr32 = ( UInt32* )_originalMethodPointer.ToPointer();
                _savedOriginalMethodAddress = *ptr32;
            }

            _dynamicMethod = CreateDynamicMethod( origMethod, _instance );
        }

        public void Hook()
        {
            unsafe
            {
                UInt32* ptr = ( UInt32* )_originalMethodPointer.ToPointer();
                *ptr = *( ( UInt32* )_hookMethodPointer.ToPointer() );
            }
        }

        public unsafe void UnHook()
        {
            UInt32* ptr = ( UInt32* )_originalMethodPointer.ToPointer();
            *ptr = _savedOriginalMethodAddress;
        }

        public void InvokeOriginalMethod( params Object[] parameters )
        {
            InternalInvokeOriginalMethod( parameters );
        }

        public T InvokeOriginalMethod<T>( params Object[] parameters )
        {
            T res = ( T )InternalInvokeOriginalMethod( parameters );
            return res;
        }

        #region - Private methods -
        private Object InternalInvokeOriginalMethod( params Object[] parameters )
        {
            var addr = _savedOriginalMethodAddress;

            Object result = null;

            var dynamicMethodParameters = new List<Object>();
            if ( _instance != null )
            {
                dynamicMethodParameters.Add( _instance );
            }

            dynamicMethodParameters.AddRange( parameters );
            dynamicMethodParameters.Add( addr );

            result = _dynamicMethod.Invoke( null, dynamicMethodParameters.ToArray() );
            return result;
        }

        private MethodInfo CreateDynamicMethod( MethodInfo method, Object instance )
        {
            var tb = CreateTypeBuilder( method );
            var proxyMethodName = method.Name + "Proxy";

            var argumentType = new List<Type>();

            if ( !method.IsStatic )
            {
                // add the this pointer as first argument
                argumentType.Add( instance.GetType() );
            }

            foreach ( var parameter in method.GetParameters() )
            {
                argumentType.Add( parameter.ParameterType );
            }

            // the address of the method to call
            argumentType.Add( typeof( UInt32 ) );

            var mb = tb.DefineMethod(
                proxyMethodName,
                MethodAttributes.Private | MethodAttributes.Static,
                method.ReturnType,
                argumentType.ToArray()
                );

            GenerateBody( method, mb.GetILGenerator(), instance );

            var proxyType = tb.CreateType();
            return proxyType.GetMethod( proxyMethodName, BindingFlags.Static | BindingFlags.NonPublic );
        }

        private static void GenerateBody( MethodInfo method, ILGenerator ilGenerator, Object instance )
        {
            // push all parameters on the stack
            var index = 0;
            var inputTypes = new List<Type>();

            if ( !method.IsStatic )
            {
                // load the instance pointer
                ilGenerator.Emit( OpCodes.Ldarg, index );
                inputTypes.Add( instance.GetType() );
                index++;
            }

            foreach ( var parameter in method.GetParameters() )
            {
                inputTypes.Add( parameter.ParameterType );
                ilGenerator.Emit( OpCodes.Ldarg, index );

                index++;
            }

            // load the function addr
            ilGenerator.Emit( OpCodes.Ldarg, index );         

            ilGenerator.EmitCalli( 
                OpCodes.Calli, 
                CallingConventions.Any, 
                method.ReturnType, 
                inputTypes.ToArray(), 
                null );

            ilGenerator.Emit( OpCodes.Ret );
        }

        private static TypeBuilder CreateTypeBuilder( MethodInfo method )
        {
            AssemblyName aName = new AssemblyName( method.Name + "DynamicAssembly" );
            AssemblyBuilder ab = AppDomain.CurrentDomain.DefineDynamicAssembly( aName, AssemblyBuilderAccess.Run );
            ModuleBuilder mb = ab.DefineDynamicModule( "dynamicModule", true );
            TypeBuilder tb = mb.DefineType( "dynamicType", TypeAttributes.Public );
            return tb;
        }

        private static IntPtr GetMethodAddress( MethodInfo method )
        {
            var methodDescriptorAddr = method.MethodHandle.Value;
            var methodAddr = new IntPtr( methodDescriptorAddr.ToInt32() + ( IntPtr.Size * 2 ) );

            return methodAddr;
        }

        private static Boolean AreMethodsSignaturesEqual( MethodBase firstMethod, MethodBase secondMethod )
        {
            var areEqual = true;

            if ( !secondMethod.CallingConvention.Equals( firstMethod.CallingConvention ) )
            {
                areEqual = false;
            }

            if ( firstMethod.IsStatic != secondMethod.IsStatic )
            {
                areEqual = false;
            }

            Type returnTypeFirstMethod = GetMethodReturnType( firstMethod );
            Type returnTypeSecondMethod = GetMethodReturnType( secondMethod );

            if ( !returnTypeFirstMethod.Equals( returnTypeSecondMethod ) )
            {
                areEqual = false;
            }

            ParameterInfo[] parametersFirstMethod = firstMethod.GetParameters();
            ParameterInfo[] parametersSecondMethod = secondMethod.GetParameters();

            if ( parametersFirstMethod.Length != parametersSecondMethod.Length )
            {
                areEqual = false;
            }

            for ( var i = 0; i < parametersFirstMethod.Length; i++ )
            {
                if ( !parametersFirstMethod[ i ].ParameterType.Equals( parametersSecondMethod[ i ].ParameterType ) )
                {
                    areEqual = false;
                }
            }

            return areEqual;
        }

        private static Type GetMethodReturnType( MethodBase method )
        {
            MethodInfo methodInfo = method as MethodInfo;

            if ( methodInfo == null )
            {
                throw new ArgumentException( "Unsupported MethodBase : " + method.GetType().Name, "method" );
            }

            return methodInfo.ReturnType;
        }
        #endregion
    }
}
